using Unity.Burst;
using Unity.Collections;
using Unity.Entities;
using Unity.Entities.Exposed;
using Unity.Entities.LowLevel.Unsafe;
using Unity.Jobs;

namespace Latios.Authoring
{
    /// <summary>
    /// A handle to a blob asset to be created by a Smart Blobber
    /// </summary>
    /// <typeparam name="TBlobType">The top of blob asset to be created</typeparam>
    public struct SmartBlobberHandle<TBlobType> where TBlobType : unmanaged
    {
        internal Entity entityWithResultBlob;
        internal bool   wasFiltered;

        /// <summary>
        /// Retrieves the blob asset after the smart blobber has run. Throws an exception if called before the smart blobber has run.
        /// </summary>
        /// <returns>The blob asset generated by the smart blobber, or Null if it was filtered or ignored</returns>
        public BlobAssetReference<TBlobType> Resolve(EntityManager entityManager)
        {
            if (wasFiltered)
                return BlobAssetReference<TBlobType>.Null;

            if (!IsValid)
                throw new System.InvalidOperationException("This handle has not been initialized.");

            if (!entityManager.HasComponent<SmartBlobberTrackingData>(entityWithResultBlob))
            {
                UnityEngine.Debug.LogError("Something internally went wrong.");
            }
            var trackingData = entityManager.GetComponentData<SmartBlobberTrackingData>(entityWithResultBlob);
            if (!trackingData.isFinalized)
                throw new System.InvalidOperationException(
                    $"The smart blobber has not processed the blob yet. Please request the blob generation prior to smart blobber execution such as during ISmartBakerData.CaptureInputsAndFilter() and do not attempt to resolve the blob until after smart blobber execution such as during SmartBaker.Process().");

            var result = entityManager.GetComponentData<SmartBlobberResult>(entityWithResultBlob);

            return result.blob.Reinterpret<TBlobType>();
        }

        /// <summary>
        /// Retrieves the blob asset after the smart blobber has run. Throws an exception if called before the smart blobber has run.
        /// </summary>
        /// <returns>The blob asset generated by the smart blobber, or Null if it was filtered or ignored</returns>
        public BlobAssetReference<TBlobType> Resolve(ref SmartBlobberResolverLookup resolverLookup)
        {
            if (wasFiltered)
                return BlobAssetReference<TBlobType>.Null;

            if (!IsValid)
                throw new System.InvalidOperationException("This handle has not been initialized.");

            if (!resolverLookup.trackingDataLookup.HasComponent(entityWithResultBlob))
            {
                UnityEngine.Debug.LogError("Something internally went wrong.");
            }
            var trackingData = resolverLookup.trackingDataLookup[entityWithResultBlob];
            if (!trackingData.isFinalized)
                throw new System.InvalidOperationException(
                    $"The smart blobber has not processed the blob yet. Please request the blob generation prior to smart blobber execution such as during ISmartBakerData.CaptureInputsAndFilter() and do not attempt to resolve the blob until after smart blobber execution such as during SmartBaker.Process().");

            var result = resolverLookup.resultLookup[entityWithResultBlob];

            return result.blob.Reinterpret<TBlobType>();
        }

        /// <summary>
        /// Returns true if this handle passed the filter stage and was potentially processed by the SmartBlobber.
        /// </summary>
        public bool IsValid => entityWithResultBlob != Entity.Null;

        public static implicit operator SmartBlobberHandleUntyped(SmartBlobberHandle<TBlobType> typed)
        {
            return new SmartBlobberHandleUntyped { entityWithResultBlob = typed.entityWithResultBlob, wasFiltered = typed.wasFiltered };
        }
    }

    /// <summary>
    /// A handle to a computed blob to be created by a smart blobber
    /// </summary>
    public struct SmartBlobberHandleUntyped
    {
        internal Entity entityWithResultBlob;
        internal bool   wasFiltered;

        /// <summary>
        /// Retrieves the blob asset after the smart blobber has run. Throws an exception if called before the smart blobber has run.
        /// </summary>
        /// <returns>The blob asset generated by the smart blobber</returns>
        public UnsafeUntypedBlobAssetReference Resolve(EntityManager entityManager)
        {
            if (wasFiltered)
                return default;

            if (!IsValid)
                throw new System.InvalidOperationException("This handle has not been initialized.");

            var trackingData = entityManager.GetComponentData<SmartBlobberTrackingData>(entityWithResultBlob);
            if (!trackingData.isFinalized)
                throw new System.InvalidOperationException(
                    $"The smart blobber has not processed the blob yet. Please request the blob generation prior to smart blobber execution such as during ISmartBakerData.CaptureInputsAndFilter() and do not attempt to resolve the blob until after smart blobber execution such as during SmartBaker.Process().");

            var result = entityManager.GetComponentData<SmartBlobberResult>(entityWithResultBlob);

            return result.blob;
        }

        /// <summary>
        /// Returns true if this handle passed the filter stage and was potentially processed by the SmartBlobber.
        /// </summary>
        public bool IsValid => entityWithResultBlob != Entity.Null;
    }

    [BurstCompile]
    public static class SmartBlobberRequestExtensions
    {
        /// <summary>
        /// Creates an additional baking-only entity to hold the SmartBlobber Request.
        /// Then invokes the Filter method of the input. If Filter() returns true,
        /// adds a SmartBlobberResult to the additional entity as well as internal
        /// components and then returns a valid SmartBlobberHandle. If Filter()
        /// returns false, returns an invalid SmartBlobberHandle.
        /// </summary>
        /// <typeparam name="TBlobType">The type of blob asset to request input for</typeparam>
        /// <typeparam name="TInputType">The type of self-filteringinput configuration</typeparam>
        /// <param name="baker">The IBaker this method extends</param>
        /// <param name="input">The input self-filtering configuration</param>
        /// <returns>A smart blobber handle that can be resolved later in the baking process</returns>
        public static SmartBlobberHandle<TBlobType> RequestCreateBlobAsset<TBlobType, TInputType>(
            this IBaker baker,
            TInputType input)
            where TBlobType : unmanaged
            where TInputType : ISmartBlobberRequestFilter<TBlobType>
        {
            var entity = baker.CreateAdditionalEntity(TransformUsageFlags.None, true);
            if (input.Filter(baker, entity))
            {
                MakeComponentTypeSet(out var typesToAdd);
                baker.AddComponent(entity, typesToAdd);
                baker.SetComponent(entity, new SmartBlobberTrackingData
                {
                    isFinalized = false
                });
                baker.AddSharedComponent(entity, new SmartBlobberBlobTypeHash { hash = BurstRuntime.GetHashCode64<TBlobType>() });
                return new SmartBlobberHandle<TBlobType> { entityWithResultBlob      = entity, wasFiltered = false };
            }
            else
            {
                return new SmartBlobberHandle<TBlobType> { entityWithResultBlob = Entity.Null, wasFiltered = true };
            }
        }

        /// <summary>
        /// Creates an additional entity to hold the Smart Blobber request and returns a SmartBlobberHandle.
        /// This method also returns the blob baking entity for the caller to add configurations to.
        /// This overload is only suited for end-user code and should not be used in this framework or other
        /// third-party library directly.
        /// </summary>
        /// <typeparam name="TBlobType">The type of blob asset to request input for</typeparam>
        /// <param name="baker">The IBaker this method extends</param>
        /// <param name="blobBakingOnlyEntity">The generated baking only entity the caller should add blob baking components to</param>
        /// <returns>A smart blobber handle that can be resolved later in the baking process</returns>
        public static SmartBlobberHandle<TBlobType> RequestCreateBlobAsset<TBlobType>(this IBaker baker, out Entity blobBakingOnlyEntity) where TBlobType: unmanaged
        {
            var entity = baker.CreateAdditionalEntity(TransformUsageFlags.None, true);
            MakeComponentTypeSet(out var typesToAdd);
            baker.AddComponent(entity, typesToAdd);
            baker.SetComponent(entity, new SmartBlobberTrackingData
            {
                isFinalized = false
            });
            baker.AddSharedComponent(entity, new SmartBlobberBlobTypeHash { hash = BurstRuntime.GetHashCode64 < TBlobType>() });
            blobBakingOnlyEntity = entity;
            return new SmartBlobberHandle<TBlobType> { entityWithResultBlob = entity, wasFiltered = false };
        }

        [BurstCompile]
        private static void MakeComponentTypeSet(out ComponentTypeSet typeSet)
        {
            typeSet = new ComponentTypeSet(ComponentType.ReadWrite<SmartBlobberResult>(), ComponentType.ReadWrite<SmartBlobberTrackingData>());
        }
    }

    /// <summary>
    /// A utility Lookup type that can be used to resolve SmartBlobberHandles in a Baking System.
    /// It can be fetched via the GetSmartBlobberResolverLookup extension methods.
    /// </summary>
    public partial struct SmartBlobberResolverLookup
    {
        [ReadOnly] internal ComponentLookup<SmartBlobberTrackingData> trackingDataLookup;
        [ReadOnly] internal ComponentLookup<SmartBlobberResult>       resultLookup;

        public void Update(ref SystemState state)
        {
            trackingDataLookup.Update(ref state);
            resultLookup.Update(ref state);
        }

        public void Update(SystemBase system)
        {
            trackingDataLookup.Update(system);
            resultLookup.Update(system);
        }
    }

    public static class SmartBlobberResolverExtensions
    {
        /// <summary>
        /// Returns an instance of SmartBlobberResolverLookup. You must call this in OnCreate().
        /// </summary>
        public static SmartBlobberResolverLookup GetSmartBlobberResolverLookup(this ref SystemState state)
        {
            return new SmartBlobberResolverLookup
            {
                trackingDataLookup = state.GetComponentLookup<SmartBlobberTrackingData>(true),
                resultLookup       = state.GetComponentLookup<SmartBlobberResult>(true)
            };
        }

        /// <summary>
        /// Returns an instance of SmartBlobberResolverLookup.
        /// </summary>
        public static SmartBlobberResolverLookup GetSmartBlobberResolverLookup(this SystemBase system)
        {
            return new SmartBlobberResolverLookup
            {
                trackingDataLookup = system.GetComponentLookup<SmartBlobberTrackingData>(true),
                resultLookup       = system.GetComponentLookup<SmartBlobberResult>(true)
            };
        }
    }

    /// <summary>
    /// Implement this interface to filter a SmartBlobber request. Request inputs should be members
    /// of the implementing type.
    /// </summary>
    /// <typeparam name="TBlobType"></typeparam>
    public interface ISmartBlobberRequestFilter<TBlobType> where TBlobType : unmanaged
    {
        /// <summary>
        /// Process a Smart Blobber request by performing input filtering and adding necessary components
        /// required by the Smart Blobber Baking Systems to the blobBakingEntity.
        /// Warning: Do not add components to the primary entity!
        /// </summary>
        /// <param name="baker">The baker to further gather inputs, add components to blobBakingEntity, and perform initial validation.</param>
        /// <param name="blobBakingEntity">The target entity for adding components to be processed by a Smart Blobber Baking System.
        /// This entity will also have a SmartBlobberResult.</param>
        /// <returns>True if the inputs are valid and further blob processing should be performed</returns>
        bool Filter(IBaker baker, Entity blobBakingEntity);
    }

    /// <summary>
    /// A component to store a generated Blob Asset. Your custom Smart Blobber should write to this component.
    /// Allocate this BlobAsset with Allocator.Persistent. Any blob assets not assigned to this component
    /// must be manually disposed.
    /// Custom Smart Blobber Baking Systems must update in SmartBlobberBakingGroup.
    /// </summary>
    [TemporaryBakingType]
    public struct SmartBlobberResult : IComponentData
    {
        public UnsafeUntypedBlobAssetReference blob;
    }

    /// <summary>
    /// A struct which is used to correctly construct Smart blobber Post-Processing systems.
    /// Such systems read the SmartBlobberResult components and perform tracking and deduplication.
    /// </summary>
    /// <typeparam name="TBlobType"></typeparam>
    public partial struct SmartBlobberTools<TBlobType> where TBlobType : unmanaged
    {
        /// <summary>
        /// Register a blob asset type for which Smart Blobber tracking should be performed.
        /// Call this in OnCreate() of one of your Smart Blobber baking systems.
        /// </summary>
        /// <param name="managedWorld">The World instance of the baking system</param>
        public void Register(World managedWorld)
        {
            if (managedWorld.GetExistingSystemManaged<Systems.SmartBlobberTypedPostProcessBakingSystem<TBlobType> >() != null)
                return;
            TypeManager.GetSystemTypeIndex(typeof(Systems.SmartBlobberTypedPostProcessBakingSystem<TBlobType>));
            var system = managedWorld.GetOrCreateSystemManaged<Systems.SmartBlobberTypedPostProcessBakingSystem<TBlobType> >();
            var group  = managedWorld.GetExistingSystemManaged<Systems.SmartBlobberCleanupBakingGroup>();
            group.AddSystemToUpdateListSafe(system.SystemHandle);
        }
    }

    [TemporaryBakingType]
    internal struct SmartBlobberTrackingData : IComponentData
    {
        public Hash128 hash;
        public Entity  thisEntity;
        public bool    isFinalized;
        public bool    isNull;
    }

    [TemporaryBakingType]
    internal struct SmartBlobberBlobTypeHash : ISharedComponentData
    {
        public long hash;
    }
}

